open Base
open OUnit2
open Helper
open Ax

module To_del = struct
  type t = State.to_del
  let to_string = function
    | `Sink(id) -> Sink_id.to_string id
    | `Src(id) -> Src_id.to_string id
  let equal : t -> t -> bool =
    Polymorphic_compare.(=)
end

let int  : int  Dyn.rtti = Dyn.create_reified ()
let bool : bool Dyn.rtti = Dyn.create_reified ()

let to_int_exn x = Option.value_exn (int#of_dyn x)
let to_bool_exn x = Option.value_exn (bool#of_dyn x)

let lift_int f d = int#to_dyn (f (to_int_exn d))
let lower_int f d = to_int_exn (f (int#to_dyn d))

;;
run_test_tt_main @@
test_list [
  ts (module Int)
    [ "math1", 3 + 4, 7
    ; "math2", 8 * 2, 16 ];
  ts (module IntOption)
    [ "opt1", Some(3), Some(1 + 2)
    ; "opt2", None,    None];
  t_log "hello_world"
    [ "hello" ; "world" ]
    (fun () -> log "hello" ; logf "w%sld" "or");
  tfails
    [ "fail_world", Failure("ello"),
      (fun () -> failwith "hello world")
    ; "math3", Division_by_zero,
      (fun () -> 3 / (2 - 4 / 2)) ];

  "Tree">:::begin
    [ (let lf = Tree.rect 5 5 in
       let () =
         Tree.(refresh_path []
                 (arm `H [ rect 1 1
                         ; arm `V [ rect 2 2
                                  ; rect 3 3
                                  ; lf ] ]))
       in
       t "refresh_path" (module List_(Int))
         [1;2]
         (Tree.path lf)) ]
  end;

  "Behavior">:::begin
    let module B = Behavior in
    let src_id1 = Src_id.create () in
    let src_id2 = Src_id.create () in
    let behav = B.(
        map (both (of_src src_id1) (of_src src_id2))
          ~f:(fun (x, y) -> (to_int_exn x, to_bool_exn y))
      ) in
    let (deps, sample) = behav.eval Var_id.empty_map in
    [
      t "deps" (module ListUnordered_(Src_id))
        [ src_id1 ; src_id2 ]
        (Set.to_list deps);
      t "sample" (module Pair_(Int)(Bool))
        (3, true)
        (sample (Map.of_alist_exn (module Src_id)
                   [ src_id1, int#to_dyn 3
                   ; src_id2, bool#to_dyn true ]));

      "of_var">:::begin
        let var = B.create_var () in
        let env_1 = B.empty_vars |> B.bind ~var ~src_id:src_id1 in
        let env_2 = B.empty_vars |> B.bind ~var ~src_id:src_id2 in
        let behav = B.(both (of_src src_id1) (of_var var)) in
        let (deps_1, _)      = behav.eval env_1 in
        let (deps_2, sample) = behav.eval env_2 in
        [
          t "deps_1" (module ListUnordered_(Src_id))
            [ src_id1 ]
            (Set.to_list deps_1);
          t "deps_2" (module ListUnordered_(Src_id))
            [ src_id1 ; src_id2 ]
            (Set.to_list deps_2);
          t "sample" (module Option_(Pair_(Int)(Bool)))
            (Some(3, true))
            (sample (Map.of_alist_exn (module Src_id)
                       [ src_id1, int#to_dyn 3
                       ; src_id2, B.upcast var true ])
             |> fun (x, y) -> Option.both (int#of_dyn x) (Some(y)));
          tfail "empty_env"
            (Failure("not in scope"))
            (fun () -> behav.eval Var_id.empty_map);
        ]
      end;
    ]
  end;

  "Comp">:::begin
    let s1 = Src_id.create () in
    let s2 = Src_id.create () in
    let s3 = Src_id.create () in
    let state = State.of_sources
                  [ s1, int#to_dyn 10
                  ; s2, int#to_dyn 20
                  ; s3, int#to_dyn 30 ] in
    let dyn_b behav ~f = Comp.dyn behav f in
    let dyn_s s ~f = dyn_b (Behavior.(map (of_src s) ~f:to_int_exn)) ~f in
    let int_fold ~init ~f g =
      Comp.fold ~init:(int#to_dyn init) ~f:(fun e -> lift_int (f e))
        (fun b -> g (Behavior.map b ~f:to_int_exn)) in
    [
      "static">:::(
        let (tree, state) =
          State.mount state @@
          Comp.arm `H [ Comp.rect 2 4
                      ; Comp.rect 9 8 ]
        in
        [
          t "tree" (module Tree)
            (Tree.arm `H [ Tree.rect 2 4 ; Tree.rect 9 8 ])
            tree;
          t "sinks" (module Bool)
            true
            (Map.is_empty state.sinks)
        ]);

      "behavior">:::(
        let sink_id = Sink_id.Private.ahead 1 in
        let (tree, state) =
          State.mount state @@
          dyn_s s1 ~f:(fun x -> Comp.rect x (x * 2))
        in
        [
          t "tree" (module Tree)
            (Tree.rect 10 20)
            tree;
          t "sinks" (module ListUnordered_(Sink_id))
            [ sink_id ]
            (Map.keys state.sinks);
          t "conns" (module MultiMap_(Src_id)(Sink_id))
            (Map.of_alist_exn (module Src_id)
               [ s1, [ sink_id ] ])
            state.conns;
        ]);

      "dependent">:::(
        (* NOTE: (implementation detail) innermost sink_ids get generated *first* *)
        let sink_id1 = Sink_id.Private.ahead 2 in
        let sink_id2 = Sink_id.Private.ahead 1 in
        let (tree, state) =
          State.mount state @@
          dyn_s s1 ~f:begin fun x ->
            if Int.(x = 10) then
              dyn_s s1 ~f:(fun x -> Comp.rect x (x * 2))
            else
              Comp.rect 99 99 end
        in
        [
          t "tree" (module Tree)
            (Tree.rect 10 20)
            tree;
          t "conns" (module MultiMap_(Src_id)(Sink_id))
            (Map.of_alist_exn (module Src_id)
               [ s1, [ sink_id1 ; sink_id2 ] ])
            state.conns;
          t "sink_to_del" (module Map_(Sink_id)(List_(To_del)))
            (Map.of_alist_exn (module Sink_id)
               [ sink_id1, [ `Sink(sink_id2) ]
               ; sink_id2, [ ] ])
            (Map.map state.sinks ~f:(fun (Sink {del;_}) -> del))
        ] @

        let (updates, state) = State.update1 state s1 (lift_int Int.succ) in
        [
          t "updates" (module List_(Tree))
            [ Tree.rect 11 22 ; Tree.rect 99 99 ]
            (List.map updates ~f:snd);
          t "updated_conns" (module MultiMap_(Src_id)(Sink_id))
            (Map.of_alist_exn (module Src_id)
               [ s1, [ sink_id1 ] ])
            state.conns;
          t "updated_to_del" (module Map_(Sink_id)(List_(To_del)))
            (Map.of_alist_exn (module Sink_id)
               [ sink_id1, [ ] ])
            (Map.map state.sinks ~f:(fun (Sink {del;_}) -> del));
          t "updated_tree" (module Tree)
            (Tree.rect 99 99)
            (let Sink {tree;_} = Map.find_exn state.sinks sink_id1 in tree);
        ]
      );

      "nested">:::(
        let sink_id1 = Sink_id.Private.ahead 3 in
        let sink_id2 = Sink_id.Private.ahead 2 in
        let sink_id3 = Sink_id.Private.ahead 1 in
        let (tree, state) =
          State.mount state @@
          dyn_s s1 ~f:begin fun x ->
            dyn_s s2 ~f:begin fun y ->
              dyn_s s3 ~f:begin fun z ->
                Comp.rect (x + z) (y + z)
              end
            end
          end
        in
        [
          t "tree" (module Tree)
            (Tree.rect 40 50)
            tree;
          t "conns" (module MultiMap_(Src_id)(Sink_id))
            (Map.of_alist_exn (module Src_id)
               [ s1, [ sink_id1 ]
               ; s2, [ sink_id2 ]
               ; s3, [ sink_id3 ] ])
            state.conns;
          t "sink_to_del" (module Map_(Sink_id)(List_(To_del)))
            (Map.of_alist_exn (module Sink_id)
               [ sink_id1, [ `Sink(sink_id2) ]
               ; sink_id2, [ `Sink(sink_id3) ]
               ; sink_id3, [ ] ])
            (Map.map state.sinks ~f:(fun (Sink {del;_}) -> del));
        ] @
        let (updates, _) =
          State.update state @@
          Map.of_alist_exn (module Src_id)
            [ s1, lift_int Int.succ
            ; s3, lift_int Int.pred ]
        in
        [
          t "updates" (module List_(Tree))
            [ Tree.rect 40 49 ]
            (List.remove_consecutive_duplicates
               (List.map updates ~f:snd)
               ~equal:Tree.equal)
        ]
      );

      "fold">:::(
        let s_new = Src_id.Private.ahead 1 in
        let (tree, state) =
          State.mount state @@
          int_fold ~init:4
            ~f:(fun (`Evt (Event.Click(x,y))) i -> i * x + y)
            (dyn_b ~f:(fun x -> Comp.rect x (x * 3)))
        in
        [
          t "tree" (module Tree)
            (Tree.rect 4 12)
            tree;
          t "src_vals" (module Map_(Src_id)(Int))
            (Map.of_alist_exn (module Src_id)
               [ s1,10 ; s2,20 ; s3,30 ; s_new,4 ])
            (Map.map state.src_vals ~f:to_int_exn);
          t "handler_effect" (module Map_(Src_id)(Int))
            (Map.of_alist_exn (module Src_id)
               [ s_new, 8899 ])
            (Map.map (tree.event_handler (Click(100, 99)))
               ~f:(fun f -> to_int_exn (f (int#to_dyn 88))));
        ]
      );

      "fold_var">:::(
        let var = Behavior.create_var () in
        let (tree, state) =
          State.mount State.empty @@
          Comp.fold_var var ~init:9 ~f:(fun _ i -> i+1) @@
          dyn_b (Behavior.of_var var)
            ~f:(fun i -> Comp.rect (i * 3) (i * 4))
        in
        [
          t "tree" (module Tree)
            (Tree.rect 27 36)
            tree;
        ] @
        let (updates, _) = State.update state (tree.event_handler (Click(0,0))) in
        [
          t "updated_tree" (module Option_(Tree))
            (Some(Tree.rect 30 40))
            (Option.map ~f:snd (List.hd updates));
        ]
      );

      "multiple_handlers">:::(
        let s1 = Src_id.Private.ahead 1 in
        let s2 = Src_id.Private.ahead 2 in
        let t_vals name exp vals =
          t name (module Option_(Pair_(Int)(Int)))
            exp
            (Option.(both
                       (Map.find vals s1 >>= int#of_dyn)
                       (Map.find vals s2 >>= int#of_dyn))) in
        let (tree, state) =
          State.mount state @@
          int_fold ~init:10 ~f:(fun _ -> (+) 1) @@ fun i ->
          int_fold ~init:20 ~f:(fun _ -> (+) 1) @@ fun j ->
          dyn_b (Behavior.both i j) ~f:(fun (i, j) -> Comp.rect i j)
        in
        let ef = tree.event_handler (Click(0,0)) in
        let (updates, _) = State.update state (tree.event_handler (Click(0,0))) in
        [
          t_vals "init_state"
            (Some(10, 20))
            state.src_vals;
          t_vals "apply_to_state"
            (Some(11, 21))
            (Src.apply_effect state.src_vals ef);
          t "updated_tree" (module Option_(Tree))
            (Some(Tree.rect 11 21))
            (Option.map ~f:snd (List.hd updates));
        ]
      );

    ]
  end;
]
